{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Tce3stUlHN0L"
      },
      "source": [
        "##### Copyright 2018 The TensorFlow Authors.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "tuOe1ymfHZPu"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "MfBg1C5NB3X0"
      },
      "source": [
        "# Use a TPU\n",
        "\n",
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://www.tensorflow.org/guide/tpu\"><img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" />View on TensorFlow.org</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/docs/blob/master/site/en/guide/tpu.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://github.com/tensorflow/docs/blob/master/site/en/guide/tpu.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://storage.googleapis.com/tensorflow_docs/docs/site/en/guide/tpu.ipynb\"><img src=\"https://www.tensorflow.org/images/download_logo_32px.png\" />Download notebook</a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Ys81cOhXOWUP"
      },
      "source": [
        "Experimental support for Cloud TPUs is currently available for Keras and Google Colab. Before you run this Colab notebooks, ensure that your hardware accelerator is a TPU by checking your notebook settings: Runtime > Change runtime type > Hardware accelerator > TPU."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "Cw0WRaChRxTL"
      },
      "outputs": [],
      "source": [
        "from __future__ import absolute_import, division, print_function, unicode_literals\n",
        "\n",
        "try:\n",
        "  # %tensorflow_version only exists in Colab.\n",
        "  %tensorflow_version 2.x\n",
        "except Exception:\n",
        "  pass\n",
        "import tensorflow as tf\n",
        "\n",
        "import os\n",
        "import tensorflow_datasets as tfds"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "_NJm-kgFO0cC"
      },
      "source": [
        "## Distribution strategies\n",
        "This guide demonstrates how to use the distribution strategy `tf.distribute.experimental.TPUStrategy` to drive a Cloud TPU and train a Keras model. A distribution strategy is an abstraction that can be used to drive models on CPU, GPUs or TPUs. Simply swap out the distribution strategy and the model will run on the given device. See the [distribution strategy guide](./distributed_training.ipynb) for more information."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "DcDPMZs-9uLJ"
      },
      "source": [
        "Below is the code that connects to a TPU and creates the `TPUStrategy` object. Note that the `tpu` argument to `TPUClusterResolver` is a special address just for Colab. In the case that you are running on Google Compute Engine (GCE), you should instead pass in the name of your CloudTPU."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "7SO23K8oRpjI"
      },
      "outputs": [],
      "source": [
        "resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='grpc://' + os.environ['COLAB_TPU_ADDR'])\n",
        "tf.config.experimental_connect_to_cluster(resolver)\n",
        "tf.tpu.experimental.initialize_tpu_system(resolver)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "gKRALGgt_kCo"
      },
      "source": [
        "Below is a simple MNIST model, unchanged from what you would use on CPU or GPU."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "DiBiN-Z_R7P7"
      },
      "outputs": [],
      "source": [
        "def create_model():\n",
        "  return tf.keras.Sequential(\n",
        "      [tf.keras.layers.Conv2D(32, 3, activation='relu', input_shape=(28, 28, 1)),\n",
        "       tf.keras.layers.Flatten(),\n",
        "       tf.keras.layers.Dense(128, activation='relu'),\n",
        "       tf.keras.layers.Dense(10)])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "qYOYjYTg_31l"
      },
      "source": [
        "## Input datasets\n",
        "Efficient use of the `tf.data.Dataset` API is critical when using a Cloud TPU, as it is impossible to use the Cloud TPUs unless you can feed them data quickly enough. See [Input Pipeline Performance Guide](./data_performance.ipynb) for details on dataset performance.\n",
        "\n",
        "For all but the simplest experimentation (using `tf.data.Dataset.from_tensor_slices` or other in-graph data) you will need to store all data files read by the Dataset in Google Cloud Storage (GCS) buckets.\n",
        "\n",
        "For most use-cases, it is recommended to convert your data into `TFRecord` format and use a `tf.data.TFRecordDataset` to read it. See [TFRecord and tf.Example tutorial](../tutorials/load_data/tfrecord.ipynb) for details on how to do this. This, however, is not a hard requirement and you can use other dataset readers (`FixedLengthRecordDataset` or `TextLineDataset`) if you prefer.\n",
        "\n",
        "Small datasets can be loaded entirely into memory using `tf.data.Dataset.cache`.\n",
        "\n",
        "Regardless of the data format used, it is strongly recommended that you use large files, on the order of 100MB. This is especially important in this networked setting as the overhead of opening a file is significantly higher.\n",
        "\n",
        "Here you should use the `tensorflow_datasets` module to get a copy of the MNIST training data. Note that `try_gcs` is specified to use a copy that is available in a public GCS bucket. If you don't specify this, the TPU will not be able to access the data that is downloaded. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "noAd416KSCo7"
      },
      "outputs": [],
      "source": [
        "def get_dataset(batch_size=200):\n",
        "  datasets, info = tfds.load(name='mnist', with_info=True, as_supervised=True,\n",
        "                             try_gcs=True)\n",
        "  mnist_train, mnist_test = datasets['train'], datasets['test']\n",
        "\n",
        "  def scale(image, label):\n",
        "    image = tf.cast(image, tf.float32)\n",
        "    image /= 255.0\n",
        "\n",
        "    return image, label\n",
        "\n",
        "  train_dataset = mnist_train.map(scale).shuffle(10000).batch(batch_size)\n",
        "  test_dataset = mnist_test.map(scale).batch(batch_size)\n",
        "\n",
        "  return train_dataset, test_dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "mgUC6A-zCMEr"
      },
      "source": [
        "## Create and train a model\n",
        "\n",
        "Nothing here is TPU specific, you would write the same code below if you had mutliple GPUs and where using a `MirroredStrategy` rather than a `TPUStrategy`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "ubmDchPqSIx0"
      },
      "outputs": [],
      "source": [
        "strategy = tf.distribute.experimental.TPUStrategy(resolver)\n",
        "with strategy.scope():\n",
        "  model = create_model()\n",
        "  model.compile(optimizer='adam',\n",
        "                loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
        "                metrics=['sparse_categorical_accuracy'])\n",
        "  \n",
        "train_dataset, test_dataset = get_dataset()\n",
        "\n",
        "model.fit(train_dataset,\n",
        "          epochs=5,\n",
        "          validation_data=test_dataset)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "0rRALBZNCO4A"
      },
      "source": [
        "## Next steps\n",
        "\n",
        "*   [Google Cloud TPU Documentation](https://cloud.google.com/tpu/docs/) - Set up and run a Google Cloud TPU.\n",
        "*   [Distributed training with TensorFlow](./distributed_training.ipynb) - How to use distribution strategy and links to many example showing best practices.\n",
        "*   [TensorFlow Official Models](https://github.com/tensorflow/models/tree/master/official) - Examples of state of the art TensorFlow 2.x models that are Cloud TPU compatible.\n",
        "*   [The Google Cloud TPU Performance Guide](https://cloud.google.com/tpu/docs/performance-guide) - Enhance Cloud TPU performance further by adjusting Cloud TPU configuration parameters for your application."
      ]
    }
  ],
  "metadata": {
    "accelerator": "TPU",
    "colab": {
      "collapsed_sections": [],
      "name": "tpu.ipynb",
      "private_outputs": true,
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
